import uvicorn

from fastapi import FastAPI, Body
from pydantic import BaseModel
from typing import Union
from bot.chatgpt.chatgpt import ChatGPTBot
from bot.baidu.wenxin import WenXinBot
from bot.aliyun.bailian import BailianBot
from bot.zhipu.zhipu import ZhipuBot

app = FastAPI()

chatgpt_bot = ChatGPTBot()
wenxin_bot = WenXinBot()
bailian_bot = BailianBot()
zhipu_bot = ZhipuBot()


class Request(BaseModel):
    bot: str
    prompt: Union[list, str]


@app.post("/api/generate")
async def generate(request: Request):
    if request.bot == "chatgpt":
        response = chatgpt_bot.generate(request.prompt)
    elif request.bot == "wenxin":
        response = wenxin_bot.ask(request.prompt)
    elif request.bot == "bailian":
        response = bailian_bot.generate(request.prompt)
    elif request.bot == "zhipu":
        response = zhipu_bot.generate(request.prompt)
    else:
        response = "Unknown bot"

    return {"response": response}


if __name__ == "__main__":
    uvicorn.run(app, host="127.0.0.1", port=8000, log_level="debug")
